from paddlers.deploy import Predictor
import numpy as np


class Classify:
    def __init__(self):
        self.predictor = Predictor("./model/classify", use_gpu=True)
        self.color = {0: [255, 0, 0],
                      1: [30, 255, 142],
                      2: [60, 0, 255],
                      3: [255, 222, 0],
                      4: [0, 0, 0]}

    def forward(self, img1):
        res = self.predictor.predict(img1)
        pred = np.zeros((res['label_map'].shape[0], res['label_map'].shape[1], 3))
        for i in range(pred.shape[0]):
            for j in range(pred.shape[1]):
                pred[i, j] = self.color[res['label_map'][i, j]]
        return pred
